from torch.utils import data 
import numpy as np
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor, Lambda, Compose
from collections import  Counter
from PIL import Image
from datasets.base.data_load import DataManager
import torch
# import matplotlib.pyplot as plt


# 1. 从内存中加载数据并转换为tensor格式的数据
# 2. 根据非独立同分布的设置，讲数据划分为多分，并返回多分数据的numpy数据数据。因为返回dataset反而不利于进一步的划分和处理。
# 主要完成数据的多个客户端划分。不划分训练数据还是测试数据。因为可能存在一下多种情况：
# a. 分为测试节点和训练节点
# b. 每个客户端既包含测试数据也包含训练数据。单个客户端还可能面临其他情况的划分。
# 3. 封装一层的目的主要是将数据集与数据集划分的程序的代码进行分离，
# a. 通过DataManager继承实现Minist划分算法。
# b. 通过load_data方法重写，加载不同的数据。
# c. 通过dataManager = DataManager启用单例模式。


class MinistLoad(data.Dataset):
    """加载数据集

    Args:
        data (加载数据): 封装成dataset
    """    
    def __init__(self, root, dataidxs=None, train=True, transform=None, target_transform=None, download=False):

        self.root = root
        self.dataidxs = dataidxs
        self.train = train
        self.transform = transform
        self.target_transform = target_transform
        self.download = download

        self.data, self.targets = self.__build_truncated_dataset__()

    def __build_truncated_dataset__(self):

        mnist_dataobj = MNIST(self.root, self.train, self.transform, self.target_transform, self.download)

        # if self.train:
        #     data = mnist_dataobj.train_data
        #     target = mnist_dataobj.train_labels
        # else:
        #     data = mnist_dataobj.test_data
        #     target = mnist_dataobj.test_labels

        data = mnist_dataobj.data
        target = mnist_dataobj.targets

        if self.dataidxs is not None:
            data = data[self.dataidxs]
            target = target[self.dataidxs]

        return data, target

    def __getitem__(self, index):
        """
        Args:
            index (int): Index

        Returns:
            tuple: (image, target) where target is index of the target class.
        """
        img, target = self.data[index], self.target[index]

        # doing this so that it is consistent with all other datasets
        # to return a PIL Image
        img = Image.fromarray(img.numpy(), mode='L')

        # print("mnist img:", img)
        # print("mnist target:", target)

        if self.transform is not None:
            img = self.transform(img)

        if self.target_transform is not None:
            target = self.target_transform(target)

        return img, target

    def __len__(self):
        return len(self.data)


class MinistDataManager(DataManager):
    # 将数据集加载到内存中。
    def __init__(self,dir='../data'):
        super(MinistDataManager,self).__init__(dir)

    # 1 加载数据、并对数据进行预处理。转换为Tensor格式
    def load_data(self):
        train_dataset = MinistLoad(self.dir, train=True, transform=ToTensor(), download=True)
        test_dataset = MinistLoad(self.dir, train=False, transform=ToTensor(),download=True)

        train_data = torch.unsqueeze(train_dataset.data,1).numpy()
        train_targets =train_dataset.targets.numpy()
        test_data = torch.unsqueeze(test_dataset.data,1).numpy()
        test_targets =test_dataset.targets.numpy()

        return train_data,train_targets,test_data,test_targets

    

# 使用单例模式加载一个数据集。在import的时候，执行这段代码，将全局变量dataManager导入到内存当中。
dataManager = MinistDataManager(dir='../data')
